﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

namespace CRFEncoder
{
    class TaggerImpl
    {
        public int feature_id_;
        public string [,] x_;
        public FeatureIndex feature_index_;
        public Node[,] node_;
        public Path[] path_cache_;
        short[] answer_;
        short[] result_;
        public short ysize_;
        public short xsize_;
        float Z_;

        public void set_node(int i, int j, int fid)
        {
            node_[i, j].lpath_cnt = 0;
            node_[i, j].rpath_cnt = 0;
            node_[i, j].fid = fid;
        }

        void forwardbackward()
        {
            if (xsize_ == 0)
            {
                return;
            }
            for (int i = 0; i < xsize_; ++i)
            {
                for (int j = 0; j < ysize_; ++j)
                {
                    Utils.calcAlpha(i, j, ref node_, ref path_cache_);
                }
            }

            for (int i = xsize_ - 1; i >= 0; --i)
            {
                for (int j = 0; j < ysize_; ++j)
                {
                    Utils.calcBeta(i, j, ref node_, ref path_cache_);
                }
            }
            Z_ = 0.0f;
            for (int j = 0; j < ysize_; ++j)
            {
                Z_ = Utils.logsumexp(Z_, node_[0, j].beta, j == 0);
            }
        }

        public int eval(int [,] merr, int [] yfreq)
        {
            int err = 0;
            for (int i = 0; i < xsize_; ++i)
            {
                if (answer_[i] != result_[i])
                {
                    ++err;
                    merr[answer_[i], result_[i]]++;
                }
                yfreq[answer_[i]]++;
            }
            return err;
        }

        public bool open(ref EncoderFeatureIndex f)
        {
            feature_index_ = f;
            ysize_ = (short)feature_index_.y_.Count;
            return true;
        }


        public bool add2(int x_num, int size, ref string[] column)
        {
            int xsize = feature_index_.xsize();
            if (size < xsize + 1)
            {
                Console.WriteLine("# x is small: size=" + size + " xsize=" + xsize);
            }

            for (int k = 0; k < size; ++k)
            {
                x_[x_num, k] = column[k];
            }

            answer_[x_num] = 0;  // dummy
            short r = ysize_;
            for (short k = 0; k < ysize_; ++k)
            {
                if (feature_index_.y(k) == column[xsize])
                {
                    r = k;
                }
            }

            if (!(r != ysize_))
            {
                Console.WriteLine("cannot find answer: " + column[xsize]);
            }

            answer_[x_num] = r;
            return true;
        }


        public bool add(string line, int x_num)
        {
            string[] column = line.Split('\t');
            int ysize = column.Length;
            if (!add2(x_num, ysize, ref column)) return false;
            return true;
        }
        public bool shrink()
        {
            feature_index_.buildFeatures(this);
            x_ = null;
            return true;
        }


        public bool read(ref StreamReader ifs)
        {
            List<string> recordList = new List<string>();
            int ysize = 0;
            while (ifs.EndOfStream == false)
            {
                string line = ifs.ReadLine();
                if (line.Length == 0 || line[0] == '\0' || line[0] == ' ' || line[0] == '\t')
                {
                    break;
                }

                string[] items = line.Split('\t');
                ysize = items.Length;
                recordList.Add(line);
            }
            xsize_ = (short)recordList.Count;
            if (recordList.Count > 0)
            {                
                x_ = new string[xsize_, ysize];
                answer_ = new short[xsize_];
                Z_ = 0.0f;

                int x_num = 0;
                foreach (string record in recordList)
                {
                    if (!add(record, x_num))
                    {
                        return false;
                    }
                    x_num++;
                }
            }

            return true;
        }


        public float gradient(ref BigArray<float> expected)
        {
            if (xsize_ == 0)
            {
                return 0.0f;
            }           
            buildLattice();
            forwardbackward();
            float s = 0.0f;

            for (int i = 0; i < xsize_; ++i)
            {
                for (int j = 0; j < ysize_; ++j)
                {
                    Utils.calcExpectation(i, j, ref expected, Z_, ysize_, ref feature_index_.feature_cache_, ref node_, ref path_cache_);
                }
            }

            for (int i = 0; i < xsize_; ++i)
            {
                short answer_val = answer_[i];
                Node answer_Node = node_[i, answer_val];
                int[] f = feature_index_.feature_cache_[answer_Node.fid];
                for (int k = 0; k < f.Length; ++k)
                {
                    --expected[f[k] + answer_val + 1];
                }
                s += answer_Node.cost;  // UNIGRAM cost

                int[] lpath_id = answer_Node.lpath_id;
                for (int j = 0; j < answer_Node.lpath_cnt; j++)
                {
                    Path lpath = path_cache_[lpath_id[j]];
                    if (lpath.lnode_y == answer_[lpath.lnode_x])
                    {
                        f = feature_index_.feature_cache_[lpath.fid];
                        int offset = lpath.lnode_y * ysize_ + lpath.rnode_y + 1;
                        for (int k = 0; k < f.Length; ++k)
                        {
                            --expected[f[k] + offset];
                        }
                        s += lpath.cost;  // BIGRAM COST
                        break;
                    }
                }
            }

            viterbi();  // call for eval()
            return Z_ - s;
        }

        public void Init(ref short[] result, ref Node[,] node, ref Path[] path_cache)
        {
            result_ = result;
            node_ = node;
            path_cache_ = path_cache;
            Array.Clear(result_, 0, result_.Length);
        }


        void viterbi()
        {
            float bestc = -1e37f;
            short best_x = -1;
            short best_y = -1;

            for (int i = 0; i < xsize_; ++i)
            {
                for (int j = 0; j < ysize_; ++j)
                {
                    bestc = -1e37f;
                    best_x = -1;
                    best_y = -1;

                    int[] lpath_id = node_[i,j].lpath_id;
                    for (int k = 0;k < node_[i,j].lpath_cnt;k++)
                    {
                        Path it = path_cache_[lpath_id[k]];
                        Node lnode = node_[it.lnode_x, it.lnode_y];
                        float cost = lnode.bestCost + it.cost + node_[i,j].cost;
                        if (cost > bestc)
                        {
                            bestc = cost;
                            best_x = it.lnode_x;
                            best_y = it.lnode_y;
                        }
                    }
                        node_[i, j].prev_x = best_x;
                        node_[i, j].prev_y = best_y;
                        node_[i, j].bestCost = (best_x >= 0 && best_y >= 0) ? bestc : node_[i, j].cost;
                }
            }

            bestc = -1e37f;
            best_x = -1;
            best_y = -1;

            short s = (short)(xsize_ - 1);
            for (short j = 0; j < ysize_; ++j)
            {
                if (bestc < node_[s,j].bestCost)
                {
                    best_x = s;
                    best_y = j;
                    bestc = node_[s,j].bestCost;
                }
            }

            if (best_x >= 0 && best_y >= 0)
            {
                for (Node n = node_[best_x, best_y];; n = node_[n.prev_x, n.prev_y])
                {
                    result_[best_x] = (short)best_y;
                    if (n.prev_x < 0 || n.prev_y < 0)
                    {
                        break;
                    }
                    best_x = n.prev_x;
                    best_y = n.prev_y;

                }
            }
        }

        public void buildLattice()
        {
            if (xsize_ == 0)
            {
                return;
            }

            feature_index_.rebuildFeatures(this);
            for (int i = 0; i < xsize_; ++i)
            {
                for (int j = 0; j < ysize_; ++j)
                {
                    feature_index_.calcCost(i, j, ref node_);
                    int[] lpath_id = node_[i,j].lpath_id;
                    for (int k = 0; k < node_[i,j].lpath_cnt; k++)
                    {
                        feature_index_.calcCost(ref path_cache_[lpath_id[k]]);
                    }
                }
            }
        }
    }
}
